# link: https://mp.weixin.qq.com/s/rFaU96l4mYzC0Kaua9jRJA

# 官方文档：https://www.backtrader.com/docu/indautoref/
'''
在编写策略时，除了常规的高开低收成交量等行情数据外，还会用到各式各样的指标（变量），
比如宏观经济指标、基本面分析指标、技术分析指标、另类数据等等。Backtrader 大致有 2 种获取指标的方式：

1、直接通过 DataFeeds 模块导入已经计算好的指标，比如《数据篇》中的导入新增指标 PE、PB；
2、在编写策略时调用 Indicators 指标模块临时计算指标，比如 5 日均线、布林带等 。
'''

# %%
import backtrader as bt
import backtrader.indicators as btind  # 导入策略分析模块
import pandas as pd
import datetime
import akshare as ak


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    # df.date = pd.to_datetime(df.date, format='%Y-%m-%d', utc=True)
    # df.set_index('date', inplace=True)
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


codes = ['sh600276', 'sh600519', 'sh603288']

# 恒瑞医药
data1 = get_data(codes[0])
# 贵州茅台
data2 = get_data(codes[1])
# 海天味业
data3 = get_data(codes[2])

# 第3节 好用的运算函数
'''
在计算指标或编写策略逻辑时，离不开算术运算、关系运算、逻辑运算、条件运算......，
为了更好的适用于Backtrader 框架的语法规则，Backtrader 的开发者还对一些常用的运算符做了优化和改进，
使用起来更简便高效：And、Or、If、All、Any、Max、Min、Sum、Cmp...

在next()中返回的结果依然是line，可以通过[num]调用各个时间节点的数值
'''


class TestStrategy(bt.Strategy):

    def __init__(self):
        self.sma5 = btind.SimpleMovingAverage(period=5)  # 5日均线
        self.sma10 = btind.SimpleMovingAverage(period=10)  # 10日均线
        # bt.And 中所有条件都满足时返回 1；有一个条件不满足就返回 0
        self.And = bt.And(self.data > self.sma5, self.data > self.sma10, self.sma5 > self.sma10)
        # bt.Or 中有一个条件满足时就返回 1；所有条件都不满足时返回 0
        self.Or = bt.Or(self.data > self.sma5, self.data > self.sma10, self.sma5 > self.sma10)
        # bt.If(a, b, c) 如果满足条件 a，就返回 b，否则返回 c
        self.If = bt.If(self.data > self.sma5, 1000, 5000)
        # bt.All,同 bt.And
        self.All = bt.All(self.data > self.sma5, self.data > self.sma10, self.sma5 > self.sma10)
        # bt.Any，同 bt.Or
        self.Any = bt.Any(self.data > self.sma5, self.data > self.sma10, self.sma5 > self.sma10)
        # bt.Max，返回同一时刻所有指标中的最大值
        self.Max = bt.Max(self.data, self.sma10, self.sma5)
        # bt.Min，返回同一时刻所有指标中的最小值
        self.Min = bt.Min(self.data, self.sma10, self.sma5)
        # bt.Sum，对同一时刻所有指标进行求和
        self.Sum = bt.Sum(self.data, self.sma10, self.sma5)
        # bt.Cmp(a,b), 如果 a>b ，返回 1；否则返回 -1
        self.Cmp = bt.Cmp(self.data, self.sma5)

    def next(self):
        print('---------- datetime', self.data.datetime.date(0), '------------------')
        print('close:', self.data[0], 'ma5:', self.sma5[0], 'ma10:', self.sma10[0])
        print('close>ma5', self.data > self.sma5, 'close>ma10', self.data > self.sma10, 'ma5>ma10',
              self.sma5 > self.sma10)
        print('self.And', self.And[0], self.data > self.sma5 and self.data > self.sma10 and self.sma5 > self.sma10)
        print('self.Or', self.Or[0], self.data > self.sma5 or self.data > self.sma10 or self.sma5 > self.sma10)
        print('self.If', self.If[0], 1000 if self.data > self.sma5 else 5000)
        print('self.All', self.All[0], self.data > self.sma5 and self.data > self.sma10 and self.sma5 > self.sma10)
        print('self.Any', self.Any[0], self.data > self.sma5 or self.data > self.sma10 or self.sma5 > self.sma10)
        print('self.Max', self.Max[0], max([self.data[0], self.sma10[0], self.sma5[0]]))
        print('self.Min', self.Min[0], min([self.data[0], self.sma10[0], self.sma5[0]]))
        print('self.Sum', self.Sum[0], sum([self.data[0], self.sma10[0], self.sma5[0]]))
        print('self.Cmp', self.Cmp[0], 1 if self.data > self.sma5 else -1)


cerebro = bt.Cerebro()
st_date = datetime.datetime(2023, 4, 1)
ed_date = datetime.datetime(2023, 5, 12)
datafeed1 = bt.feeds.PandasData(dataname=data1, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed1, name=codes[0])
cerebro.addstrategy(TestStrategy)
result = cerebro.run()
